# —— 顶部：第三方日志静音，一次到位 ——
import os, warnings, logging

# 先用环境变量把常见库的日志降到最低（需要在 import cv2 / paddleocr 之前）
os.environ.setdefault("OPENCV_LOG_LEVEL", "SILENT")   # OpenCV
os.environ.setdefault("GLOG_minloglevel", "2")        # 一些底层库（Paddle 依赖）
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")    # 若本机装了 TF/ONNX 之类

warnings.filterwarnings("ignore")                     # Python warnings 也关掉

# OpenCV 运行时再强压一遍
import cv2
try:
    # OpenCV 4.5+ 推荐
    cv2.utils.logging.setLogLevel(cv2.utils.logging.LOG_LEVEL_SILENT)
except Exception:
    try:
        # 某些版本是这个 API
        cv2.setLogLevel(cv2.LOG_LEVEL_SILENT)
    except Exception:
        pass

# PaddleOCR 的日志降级到 ERROR（只在出错时才打印）
logging.getLogger("ppocr").setLevel(logging.ERROR)
logging.getLogger("ppocr").propagate = False

# 如果你初始化 OCR，请把 show_log 也关掉：
from paddleocr import PaddleOCR
ocr = PaddleOCR(use_angle_cls=True, lang='en', show_log=False)

# -*- coding: utf-8 -*-
"""
OCR 快速验证 + 结果落盘（自动去重 + 时间戳）
- 扫描 D:\tennis\代码与模型数据\data\images 的整帧图片
- 左下角 ROI 做 OCR，未识别到数字则丢弃
- 相邻重复结果稳定 N 帧才记一次
- 结果写到 D:\tennis\代码与模型数据\outputs\ocr_results.csv
"""

import os
import re
import csv
import numpy as np
import cv2
from paddleocr import PaddleOCR

# --------- 关闭 OpenCV 冗余日志（红色 WARN） ----------
try:
    cv2.utils.logging.setLogLevel(cv2.utils.logging.LOG_LEVEL_ERROR)
except Exception:
    os.environ["OPENCV_LOG_LEVEL"] = "ERROR"

# --------- 路径与参数（按需改） ----------
FRAMES_DIR = r"D:\tennis\代码与模型数据\data\images"   # 整帧图片目录（不在 code 里）
OUT_DIR     = r"D:\tennis\代码与模型数据\outputs"
OUT_CSV     = os.path.join(OUT_DIR, "ocr_results.csv")

FPS = 25             # 你的帧率
STABLE_N = 3         # 去重阈值：稳定出现 N 帧才记一次

# 左下角比分牌 ROI（百分比，适配 1920x1080 的这场样式）
ROI_PERCENT = dict(x=0.06, y=0.79, w=0.26, h=0.18)

# --------- OCR 实例 ----------
ocr = PaddleOCR(use_angle_cls=True, lang="en")

def safe_imread(path: str):
    """优先 imread，失败用 imdecode 兜底（兼容缺编解码的环境）"""
    img = cv2.imread(path)
    if img is None:
        try:
            buf = np.fromfile(path, dtype=np.uint8)
            img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
        except Exception:
            img = None
    return img

def cut_roi(image: np.ndarray) -> np.ndarray:
    h, w = image.shape[:2]
    x = int(w * ROI_PERCENT["x"])
    y = int(h * ROI_PERCENT["y"])
    ww = int(w * ROI_PERCENT["w"])
    hh = int(h * ROI_PERCENT["h"])
    x = max(0, min(x, w - 1))
    y = max(0, min(y, h - 1))
    ww = max(1, min(ww, w - x))
    hh = max(1, min(hh, h - y))
    return image[y:y+hh, x:x+ww]

def fname_to_timestamp(fname: str) -> tuple[str, float]:
    """frame_0598.jpg -> ([00:23], 23.92)"""
    m = re.search(r"(\d+)", fname)
    if not m:
        return ("[--:--]", 0.0)
    idx = int(m.group(1))
    sec = idx / FPS
    mm = int(sec // 60)
    ss = int(sec % 60)
    return (f"[{mm:02d}:{ss:02d}]", sec)

def has_any_digit(ocr_lines) -> bool:
    """只要任意文本里出现数字就算含比分信息"""
    if not ocr_lines:
        return False
    for line in ocr_lines:
        if not line or not line[1]:
            continue
        txt = str(line[1][0])
        if any(ch.isdigit() for ch in txt):
            return True
    return False

def lines_to_display(ocr_lines):
    out = []
    for line in ocr_lines or []:
        if not line or not line[1]:
            continue
        text, conf = line[1]
        out.append((str(text), float(conf)))
    return out

def main():
    if not os.path.isdir(FRAMES_DIR):
        print(f"[ERROR] 找不到图片目录：{FRAMES_DIR}")
        return
    os.makedirs(OUT_DIR, exist_ok=True)

    # 收集图片
    files = [f for f in os.listdir(FRAMES_DIR)
             if f.lower().endswith((".jpg", ".jpeg", ".png"))]

    # 统一排序键：返回 (是否有数字, 数字键, 原名)
    def key_fn(f):
        m = re.search(r"(\d+)", f)
        if m:
            return (0, int(m.group(1)), f)  # 先按数字排序
        return (1, float("inf"), f)        # 无数字的排在后面，但不会和 int 混比

    files.sort(key=key_fn)

    kept, dropped = 0, 0
    last_key = None
    stable_cnt = 0
    rows = []

    for fname in files:
        fpath = os.path.join(FRAMES_DIR, fname)

        img = safe_imread(fpath)
        if img is None:
            print(f"[无效图片] 跳过: {fname} -> OpenCV 无法解码")
            dropped += 1
            stable_cnt = 0
            last_key = None
            continue

        roi = cut_roi(img)

        result = ocr.ocr(roi, cls=True)
        ocr_lines = result[0] if result else None

        if not has_any_digit(ocr_lines):
            print(f"[无数字] 丢弃: {fname}")
            dropped += 1
            stable_cnt = 0
            last_key = None
            continue

        # 关键串（仅保留字母数字，用于相邻去重）
        kv = "".join(
            re.sub(r"[^A-Za-z0-9]+", "", (ln[1][0] if ln and ln[1] else ""))
            for ln in (ocr_lines or [])
        )

        if kv == last_key:
            stable_cnt += 1
        else:
            last_key = kv
            stable_cnt = 1

        if stable_cnt >= STABLE_N:
            timestamp_str, sec = fname_to_timestamp(fname)
            pairs = lines_to_display(ocr_lines)

            print(f"{timestamp_str} 识别自: {fname}")
            for t, c in pairs:
                print(f" - {t:<8} (conf: {c:.2f})")

            all_text = " | ".join([p[0] for p in pairs])
            digits   = " ".join(re.findall(r"\d+", all_text))
            rows.append([timestamp_str, f"{sec:.2f}", fname, all_text, digits])
            kept += 1

            # 记一次后清零等待下次新结果
            stable_cnt = 0
            last_key   = None

    # 写 CSV
    with open(OUT_CSV, "w", newline="", encoding="utf-8-sig") as f:
        writer = csv.writer(f)
        writer.writerow(["timestamp", "sec", "frame", "text", "digits"])
        writer.writerows(rows)

    print(f"\n[INFO] 完成: 已写入 {len(rows)} 条去重结果 -> {OUT_CSV}")
    print(f"[INFO] 统计: 保留 {kept} 条，丢弃 {dropped} 帧（无数字/读图失败等）")

if __name__ == "__main__":
    main()
